# -*- coding: utf-8 -*-
# pathon version > 3.
import os


class FilesList(object):
    def __init__(self, path):
        self.path = path
        # 指定路径下的文件列表
        self.files_list = []
        for file in os.listdir(path):
            self.files_list.append(os.path.splitext(file)[0])
        # 文件数量
        self.len = len(self.files_list)

    def split(self, split=(1, 0, 0)):
        '''
        split = (train:val:test)
        输入百分比，和为1，如（0.8, 0.1, 0.1）
        '''
        # 获取分割后的数据集的样本数量
        train = round(self.len*split[0])  # 四舍五入
        val = round(self.len*split[1])
        test = self.len - train - val

        # 获取分割后的数据集的样本文件
        train_files = self.files_list[:train]
        val_files = self.files_list[train:-test]
        test_files = self.files_list[-test:]

        return (train_files, val_files, test_files)


def split_files():
    '''
    根据文件名整理文件，将文件复制到指定目录
    '''
    for i in os.listdir("./segment"):
        if i.startswith("00002_") or i.startswith("00003_") or i.startswith("00005_") or i.startswith("00007_") or i.startswith("00009_"):
            os.system("cp ./segment/" + i + " ./dfb/100")
        if i.startswith("00001_") or i.startswith("00004_") or i.startswith("00006_") or i.startswith("00008_"):
            os.system("cp ./segment/" + i + " ./dfb/200")


if __name__ == "__main__":

    xyz_path = "./src_points"
    label_path = "./points_label"
    xyz = FilesList(xyz_path)
    a = xyz.split((0.6, 0.2, 0.2))
    # print(xyz.files_list)

    # 文件后缀名
    data_ext = ".pts"
    label_ext = ".seg"

    # 将data文件复制到相应文件夹
    os.mkdir("train_data")
    os.mkdir("val_data")
    os.mkdir("test_data")
    for train in a[0]:
        train = train+ data_ext
        os.system("cp " + os.path.join(xyz_path, train) +
                  " " + os.path.join("train_data", train))
    for val in a[1]:
        val = val+ data_ext
        os.system("cp " + os.path.join(xyz_path, val) +
                  " " + os.path.join("val_data", val))
    for test in a[2]:
        test = test+ data_ext
        os.system("cp " + os.path.join(xyz_path, test) +
                  " " + os.path.join("test_data", test))

    # 将label文件复制到相应文件夹
    os.mkdir("train_label")
    os.mkdir("val_label")
    os.mkdir("test_label")
    # 遍历列表中的文件名
    for train in a[0]:
        train = train+ label_ext  # 为文件名添加后缀
        # 复制文件
        os.system("cp " + os.path.join(label_path, train) +
                  " " + os.path.join("train_label", train))
    for val in a[1]:
        val = val+ label_ext
        os.system("cp " + os.path.join(label_path, val) +
                  " " + os.path.join("val_label", val))
    for test in a[2]:
        test = test+ label_ext
        os.system("cp " + os.path.join(label_path, test) +
                  " " + os.path.join("test_label", test))
